#include "geometry_msgs/msg/pose_array.hpp"
#include <cmath>
#include <opencv2/opencv.hpp>

#include "final_exam/armor.h"
#include "final_exam/param.h"
#include "final_exam/armor_detect.h"

void ArmorDetect::BGR2Gray(cv::Mat &src, cv::Mat &gray) {
  cv::Mat channels[3];
  cv::split(src, channels);
  if (this->enemy == param::RED) {
    gray = channels[2] - channels[0];
  } else {
    gray = channels[0] - channels[2];
  }
}

void ArmorDetect::imageEnhance(cv::Mat &gray, cv::Mat &enhance) {
  // 二值化
  cv::Mat binary_image;
  cv::threshold(gray, binary_image, param::binary_thres, 255,
                cv::THRESH_BINARY);

  // 膨胀
  cv::Mat tmp_image;
  static cv::Mat element;
  element = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(3, 3));
  cv::dilate(binary_image, tmp_image, element);
  // 腐蚀
  element = cv::getStructuringElement(cv::MORPH_RECT, cv::Size(3, 3));
  cv::erode(tmp_image, enhance, element);
  // 涂黑左上侧区域，减少干扰
  cv::Rect roi(0, 0, enhance.cols / 4, enhance.rows / 4);
  enhance(roi) = cv::Scalar(0);
  // cv::imshow("gray_img", enhance);
}

void ArmorDetect::findLightBar(cv::Mat &enhance,
                               std::vector<cv::RotatedRect> &light_bars) {
  std::vector<std::vector<cv::Point>> contours;
  cv::findContours(enhance, contours, cv::RETR_EXTERNAL,
                   cv::CHAIN_APPROX_SIMPLE);
  for (const auto &contour : contours) {
    if (contour.size() < 5) // 拟合椭圆至少需要 5 个点
      continue;

    cv::RotatedRect box = cv::fitEllipse(contour);
    double aspect_ratio = box.size.height * 1.0 / box.size.width;
    if (aspect_ratio < 1.0)
      aspect_ratio = 1.0 / aspect_ratio;

    if (aspect_ratio > param::max_light_ratio ||
        aspect_ratio < param::min_light_ratio)
      continue; // 长宽比筛选

    if (box.size.area() < param::min_light_area ||
        box.size.area() > param::max_light_area)
      continue; // 面积筛选

    // // 1. 获取矩形的四个角点
    // cv::Point2f vertices[4];
    // box.points(vertices);

    // // 2. 依次连接四个角点绘制四条边
    // for (int i = 0; i < 4; i++) {
    //   // 使用 cv::line 函数连接点 i 和点 (i+1)%4
    //   cv::line(this->img, vertices[i], vertices[(i + 1) % 4],
    //            cv::Scalar(0, 255, 255), 2); // 颜色：黄色 (BGR: 蓝0, 绿255, 红255)
    // }

    // std::string text = std::to_string(box.size.area()).substr(0, 5);
    // cv::putText(this->img, text, vertices[0] + cv::Point2f(-10, -20),
    //             cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255, 255, 255), 1);
    light_bars.push_back(box);
  }
}

void ArmorDetect::matchLightBar(std::vector<cv::RotatedRect> &light_bars,
                                std::vector<Armor> &armors) {
  for (size_t i = 0; i < light_bars.size(); ++i) {
    for (size_t j = i + 1; j < light_bars.size(); ++j) {
      const auto &bar1 = light_bars[i];
      const auto &bar2 = light_bars[j];
      // 计算两个灯条的中心点距离
      double center_distance = cv::norm(bar1.center - bar2.center);
      // 计算两个灯条的平均高度
      double avg_height = (bar1.size.height + bar2.size.height) / 2.0;
      double height_diff = std::max(bar1.size.height, bar2.size.height) * 1.0 /
                           std::min(bar1.size.height, bar2.size.height);
      // 高度差筛选
      if (height_diff > param::max_armor_height_diff)
        continue;

      // 计算两个灯条的角度差
      double angle_diff = std::abs(bar1.angle - bar2.angle);
      if (angle_diff > param::max_angle_diff)
        continue; // 角度差筛选

      // 计算长宽比
      double length_ratio = center_distance / avg_height;
      if (length_ratio > param::max_armor_ratio ||
          length_ratio < param::min_armor_ratio)
        continue; // 长宽比筛选

      double area = center_distance * avg_height;
      if (area > param::max_area || area < param::min_area ||
          isnan(area)) // 注意判断 nan
        continue;      // 面积筛选

      // 获取装甲板图像
      cv::Mat armor_img;
      this->getArmorImg(this->src, {bar1, bar2}, armor_img);
      // if (type == "Unknown")
      //   continue; // 兵种识别筛选

      armors.push_back(Armor(bar1, bar2));
      // armors.push_back(Armor(bar1, bar2));
      // std::cerr << area << std::endl;
    }
  }
}

void ArmorDetect::drawArmor() {
  // 显示文字信息，存入位置消息
  auto putPoseMsg = [this](const Armor &armor) -> void {
    auto &[bar1, bar2, type, confidence_level, tvec, rvec] = armor;
    // 使用 rvec 直接计算四元数并存入 Pose（rpy 的计算已移至其他节点）
    geometry_msgs::msg::Pose pose_msg = this->poseFromRt(tvec, rvec);
    this->armor_poses_msg.poses.push_back(pose_msg);
  };

  // 可预测的装甲板输出
  for (auto [predictor, armor] : this->predictors) {
    auto &[bar1, bar2, type, confidence_level, tvec, rvec] = armor;
    cv::Point3f predicted_pt = predictor.getState();
    cv::Mat tvec_predicted = (cv::Mat_<double>(3, 1) << predicted_pt.x,
                              predicted_pt.y, predicted_pt.z);
    std::vector<cv::Point2f> image_points =
        this->pose_solver.project(rvec, tvec);
    for (size_t i = 0; i < image_points.size(); ++i) {
      cv::line(this->img, image_points[i],
               image_points[(i + 1) % image_points.size()],
               cv::Scalar(0, 255, 0), 2);
    }

    putPoseMsg({bar1, bar2, type, confidence_level, tvec_predicted, rvec});
  }
}
// 将 tvec (3x1) 和 rvec (3x1) 转换为 geometry_msgs::msg::Pose
geometry_msgs::msg::Pose ArmorDetect::poseFromRt(const cv::Mat &tvec,
                                                 const cv::Mat &rvec) {
  geometry_msgs::msg::Pose pose_msg;
  // 位置
  pose_msg.position.x = tvec.at<double>(0);
  pose_msg.position.y = tvec.at<double>(1);
  pose_msg.position.z = tvec.at<double>(2);

  // 旋转：rvec -> R -> quaternion
  cv::Mat R;
  cv::Rodrigues(rvec, R);
  double m00 = R.at<double>(0, 0);
  double m01 = R.at<double>(0, 1);
  double m02 = R.at<double>(0, 2);
  double m10 = R.at<double>(1, 0);
  double m11 = R.at<double>(1, 1);
  double m12 = R.at<double>(1, 2);
  double m20 = R.at<double>(2, 0);
  double m21 = R.at<double>(2, 1);
  double m22 = R.at<double>(2, 2);

  double qw, qx, qy, qz;
  double trace = m00 + m11 + m22;
  if (trace > 0.0) {
    double s = std::sqrt(trace + 1.0) * 2.0; // s = 4*qw
    qw = 0.25 * s;
    qx = (m21 - m12) / s;
    qy = (m02 - m20) / s;
    qz = (m10 - m01) / s;
  } else if (m00 > m11 && m00 > m22) {
    double s = std::sqrt(1.0 + m00 - m11 - m22) * 2.0; // s = 4*qx
    qw = (m21 - m12) / s;
    qx = 0.25 * s;
    qy = (m01 + m10) / s;
    qz = (m02 + m20) / s;
  } else if (m11 > m22) {
    double s = std::sqrt(1.0 + m11 - m00 - m22) * 2.0; // s = 4*qy
    qw = (m02 - m20) / s;
    qx = (m01 + m10) / s;
    qy = 0.25 * s;
    qz = (m12 + m21) / s;
  } else {
    double s = std::sqrt(1.0 + m22 - m00 - m11) * 2.0; // s = 4*qz
    qw = (m10 - m01) / s;
    qx = (m02 + m20) / s;
    qy = (m12 + m21) / s;
    qz = 0.25 * s;
  }

  pose_msg.orientation.x = qx;
  pose_msg.orientation.y = qy;
  pose_msg.orientation.z = qz;
  pose_msg.orientation.w = qw;

  return pose_msg;
}

void ArmorDetect::reorderArmorPoints(
    const std::vector<cv::Point2f> &input_points,
    std::vector<cv::Point2f> &output_points) {
  // 1. 复制输入点，准备排序
  std::vector<cv::Point2f> sorted_x = input_points;

  // 2. 按 X 坐标从小到大排序
  std::sort(
      sorted_x.begin(), sorted_x.end(),
      [](const cv::Point2f &a, const cv::Point2f &b) { return a.x < b.x; });

  // Pts[0] 和 Pts[1] 是左边的两个点
  // Pts[2] 和 Pts[3] 是右边的两个点

  // 3. 在左边的点中，按 Y 坐标排序（Y 较小的是上，Y 较大的是下）
  cv::Point2f p_LU, p_LD; // 左上，左下
  if (sorted_x[0].y < sorted_x[1].y) {
    p_LU = sorted_x[0];
    p_LD = sorted_x[1];
  } else {
    p_LU = sorted_x[1];
    p_LD = sorted_x[0];
  }

  // 4. 在右边的点中，按 Y 坐标排序
  cv::Point2f p_RU, p_RD; // 右上，右下
  if (sorted_x[2].y < sorted_x[3].y) {
    p_RU = sorted_x[2];
    p_RD = sorted_x[3];
  } else {
    p_RU = sorted_x[3];
    p_RD = sorted_x[2];
  }

  // 5. 按照 3D 模型顺序 [LU, LD, RD, RU] 赋值给输出向量
  output_points.resize(4);
  output_points[0] = p_LU; // 你的 3D 模型点 0
  output_points[1] = p_LD; // 你的 3D 模型点 1
  output_points[2] = p_RD; // 你的 3D 模型点 2
  output_points[3] = p_RU; // 你的 3D 模型点 3
}

void ArmorDetect::solveArmorPose(std::vector<Armor> &armors) {
  for (auto &[bar1, bar2, type, confidence_level, tvec, rvec] : armors) {
    // 解算位姿
    std::vector<cv::Point2f> raw_image_points(4);
    std::vector<cv::Point2f> image_points(4);
    cv::RotatedRect armor_rect;
    armor_rect.center = (bar1.center + bar2.center) / 2.0;
    double center_distance = cv::norm(bar1.center - bar2.center);
    double avg_height = (bar1.size.height + bar2.size.height) / 2.0;
    armor_rect.size = cv::Size2f(center_distance, avg_height);
    armor_rect.angle = (bar1.angle + bar2.angle) / 2.0;
    armor_rect.points(raw_image_points.data());
    reorderArmorPoints(raw_image_points, image_points);

    this->pose_solver.solve(image_points);
    rvec = this->pose_solver.getRvec();
    tvec = this->pose_solver.getTvec();
    cv::Point3f tvec_point(tvec);
    this->kalmanCorrect(tvec_point,
                        {bar1, bar2, type, confidence_level, tvec, rvec});
    tvec = (cv::Mat_<double>(3, 1) << tvec_point.x, tvec_point.y, tvec_point.z);
    // std::cerr << "image:\n" << image_points << std::endl;
    // std::cerr << "solve:\n" << this->pose_solver.project(rvec, tvec) <<
    // std::endl;
  }
}

void ArmorDetect::getArmorImg(const cv::Mat &src, const Armor &armor,
                              cv::Mat &armor_img) {
  cv::Point2f vertices[4];
  cv::RotatedRect armor_rect;
  const auto &[bar1, bar2, type, confidence_level, tvec, rvec] = armor;
  armor_rect.center = (bar1.center + bar2.center) / 2.0;
  double center_distance =
      cv::norm(bar1.center - bar2.center) * param::width_zoom_factor;

  // 高度缩放，得到整个装甲板
  double avg_height =
      (bar1.size.height + bar2.size.height) / 2.0 * param::length_zoom_factor;
  armor_rect.size = cv::Size2f(center_distance, avg_height);
  armor_rect.angle = (bar1.angle + bar2.angle) / 2.0;
  armor_rect.points(vertices);
  cv::Point2f target_vertices[4] = {
      cv::Point2f(0, param::target_height - 1),
      cv::Point2f(param::target_width - 1, param::target_height - 1),
      cv::Point2f(param::target_width - 1, 0), cv::Point2f(0, 0)};

  // 透视变换
  cv::Mat M = cv::getPerspectiveTransform(vertices, target_vertices);
  cv::warpPerspective(src, armor_img, M,
                      cv::Size(param::target_width, param::target_height));
}

void ArmorDetect::kalmanPredict() {
  for (auto it = this->predictors.begin(); it != this->predictors.end();) {
    (void)it->first.predict();
    if (it->first.getMissingCount() > param::max_missing_count) {
      it = this->predictors.erase(it);
    } else {
      ++it;
    }
  }
}

void ArmorDetect::kalmanCorrect(cv::Point3f &pt, const Armor &armor) {
  for (auto &predictor : this->predictors) {
    double distance = cv::norm(predictor.first.getState() - pt);
    if (distance < param::max_predict_distance) {
      pt = predictor.first.correct(pt);
      predictor.second = armor; // 更新装甲板信息
      return;
    }
  }

  Predictor new_predictor(param::delta_t);
  new_predictor.initKalmanFilter(pt);
  this->predictors.push_back({new_predictor, armor});
}

geometry_msgs::msg::PoseArray ArmorDetect::getArmorPosesMsg() {
  return this->armor_poses_msg;
}
void ArmorDetect::setEnemy(bool enemy) { this->enemy = enemy; }

void ArmorDetect::detect(cv::Mat &src, cv::Mat &dst) {
  // 清空位置消息
  armor_poses_msg.poses.clear();

  cv::Mat gray, enhance;
  this->img = src;
  this->src = src.clone();
  this->BGR2Gray(src, gray);
  this->imageEnhance(gray, enhance);

  std::vector<cv::RotatedRect> light_bars;
  this->findLightBar(enhance, light_bars);

  std::vector<Armor> armors;
  this->matchLightBar(light_bars, armors);
  this->kalmanPredict();
  this->solveArmorPose(armors);
  this->drawArmor();

  dst = this->img;
}

void ArmorDetect::setCameraParams(const cv::Mat &camera_matrix,
                                  const cv::Mat &dist_coeffs) {
  this->pose_solver.setCameraParams(camera_matrix, dist_coeffs);
}

void ArmorDetect::setObjectPoints(
    const std::vector<cv::Point3f> &object_points) {
  this->pose_solver.setObjectPoints(object_points);
}